diff --git a/tests/test_template.py b/tests/test_template.py index 102df3b56..ed22bae3d 100644 --- a/tests/test_template.py +++ b/tests/test_template.py @@ -3,9 +3,10 @@ from pydantic import BaseModel import pytest import re import importlib -from typing import Dict +from typing import Dict, List, Optional from langflow.utils.util import ( build_template_from_class, + build_template_from_function, format_dict, get_base_classes, get_default_factory, @@ -26,6 +27,54 @@ class Child(Parent): child_field: int +class ExampleClass1(BaseModel): + """Example class 1.""" + + def __init__(self, data: Optional[List[int]] = None): + self.data = data or [1, 2, 3] + + +class ExampleClass2(BaseModel): + """Example class 2.""" + + def __init__(self, data: Optional[Dict[str, int]] = None): + self.data = data or {"a": 1, "b": 2, "c": 3} + + +def example_loader_1() -> ExampleClass1: + """Example loader function 1.""" + return ExampleClass1() + + +def example_loader_2() -> ExampleClass2: + """Example loader function 2.""" + return ExampleClass2() + + +def test_build_template_from_function(): + type_to_loader_dict = { + "example1": example_loader_1, + "example2": example_loader_2, + } + + # Test with valid name + result = build_template_from_function("ExampleClass1", type_to_loader_dict) + + assert "template" in result + assert "description" in result + assert "base_classes" in result + + # Test with add_function=True + result_with_function = build_template_from_function( + "ExampleClass1", type_to_loader_dict, add_function=True + ) + assert "function" in result_with_function["base_classes"] + + # Test with invalid name + with pytest.raises(ValueError, match=r".* not found"): + build_template_from_function("NonExistent", type_to_loader_dict) + + # Test build_template_from_class def test_build_template_from_class(): type_to_cls_dict: Dict[str, type] = {"parent": Parent, "child": Child}