wehub-resource-sync
90 lines
2.5 KiB
Python
90 lines
2.5 KiB
Python
from dataclasses import asdict, dataclass
|
|
from typing import Optional
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
from mlflow.models.rag_signatures import ChatCompletionRequest
|
|
from mlflow.pyfunc.utils.input_converter import _hydrate_dataclass
|
|
|
|
|
|
def test_hydrate_dataclass_input_no_dataclass():
|
|
# Define a class that is not a dataclass
|
|
class NotADataclass:
|
|
pass
|
|
|
|
# Create some dummy data as a pandas df
|
|
data = {"a": 1, "b": 2}
|
|
df = pd.DataFrame(data, index=[0])
|
|
|
|
# Check that an error is raised when trying to hydrate the dataclass
|
|
with pytest.raises(ValueError, match="NotADataclass is not a dataclass"):
|
|
_hydrate_dataclass(NotADataclass, df.iloc[0])
|
|
|
|
|
|
def test_hydrate_dataclass_simple():
|
|
# Define a dataclass
|
|
@dataclass
|
|
class MyDataclass:
|
|
a: int
|
|
b: int
|
|
|
|
# Create some dummy data as a pandas df
|
|
df = pd.DataFrame({"a": [1], "b": [2]})
|
|
|
|
# Check that the dataclass is hydrated
|
|
result = _hydrate_dataclass(MyDataclass, df.iloc[0])
|
|
assert result == MyDataclass(a=1, b=2)
|
|
|
|
|
|
def test_hydrate_dataclass_complex():
|
|
# Define a more complex dataclass
|
|
@dataclass
|
|
class MyDataclass:
|
|
a: int
|
|
b: int
|
|
|
|
@dataclass
|
|
class MyListDataclass:
|
|
c: list[MyDataclass]
|
|
|
|
# Create some dummy data as a pandas df
|
|
df = pd.DataFrame({"c": [[{"a": 1, "b": 2}, {"a": 3, "b": 4}]]})
|
|
|
|
# Check that the dataclass is hydrated
|
|
result = _hydrate_dataclass(MyListDataclass, df.iloc[0])
|
|
assert result == MyListDataclass(c=[MyDataclass(a=1, b=2), MyDataclass(a=3, b=4)])
|
|
|
|
|
|
@dataclass
|
|
class CustomInput:
|
|
id: int = 0
|
|
|
|
|
|
@dataclass
|
|
class FlexibleChatCompletionRequest(ChatCompletionRequest):
|
|
custom_input: Optional[CustomInput] = None # noqa: UP045
|
|
another_custom_input: CustomInput | None = None
|
|
|
|
|
|
def test_hydrate_child_dataclass():
|
|
result = _hydrate_dataclass(
|
|
FlexibleChatCompletionRequest,
|
|
asdict(
|
|
FlexibleChatCompletionRequest(
|
|
custom_input=CustomInput(), another_custom_input=CustomInput()
|
|
)
|
|
),
|
|
)
|
|
assert result == FlexibleChatCompletionRequest(
|
|
custom_input=CustomInput(), another_custom_input=CustomInput()
|
|
)
|
|
|
|
|
|
def test_hydrate_optional_dataclass():
|
|
result = _hydrate_dataclass(
|
|
FlexibleChatCompletionRequest,
|
|
asdict(FlexibleChatCompletionRequest(custom_input=None, another_custom_input=None)),
|
|
)
|
|
assert result == FlexibleChatCompletionRequest(custom_input=None, another_custom_input=None)
|