test_beta_decorator.py — langchain Source File

Source Code

import inspect
import warnings
from typing import Any

import pytest
from pydantic import BaseModel

from langchain_core._api.beta_decorator import beta, warn_beta


@pytest.mark.parametrize(
    ("kwargs", "expected_message"),
    [
        (
            {
                "name": "OldClass",
                "obj_type": "class",
            },
            "The class `OldClass` is in beta. It is actively being worked on, so the "
            "API may change.",
        ),
        (
            {
                "message": "This is a custom message",
                "name": "FunctionA",
                "obj_type": "",
                "addendum": "",
            },
            "This is a custom message",
        ),
        (
            {
                "message": "",
                "name": "SomeFunction",
                "obj_type": "",
                "addendum": "Please migrate your code.",
            },
            "`SomeFunction` is in beta. It is actively being worked on, "
            "so the API may "
            "change. Please migrate your code.",
        ),
    ],
)
def test_warn_beta(kwargs: dict[str, Any], expected_message: str) -> None:
    """Test warn beta."""
    with warnings.catch_warnings(record=True) as warning_list:
        warnings.simplefilter("always")

        warn_beta(**kwargs)

        assert len(warning_list) == 1
        warning = warning_list[0].message
        assert str(warning) == expected_message


@beta()
def beta_function() -> str:
    """Original doc."""
    return "This is a beta function."

// ... (319 more lines)