Extend base jinja2 extension with hass requirement and tests (home-assistant#156403)

Author: frenck

homeassistant/helpers/template/extensions/base.py

@@ -11,6 +11,7 @@
 from jinja2.parser import Parser
 
 if TYPE_CHECKING:
+    from homeassistant.core import HomeAssistant
     from homeassistant.helpers.template import TemplateEnvironment
 
 
@@ -26,6 +27,7 @@ class TemplateFunction:
     limited_ok: bool = (
         True  # Whether this function is available in limited environments
     )
+    requires_hass: bool = False  # Whether this function requires hass to be available
 
 
 class BaseTemplateExtension(Extension):
@@ -44,6 +46,10 @@ def __init__(
 
         if functions:
             for template_func in functions:
+                # Skip functions that require hass when hass is not available
+                if template_func.requires_hass and self.environment.hass is None:
+                    continue
+
                 # Skip functions not allowed in limited environments
                 if self.environment.limited and not template_func.limited_ok:
                     continue
@@ -55,6 +61,24 @@ def __init__(
                 if template_func.as_test:
                     environment.tests[template_func.name] = template_func.func
 
+    @property
+    def hass(self) -> HomeAssistant:
+        """Return the Home Assistant instance.
+
+        This property should only be used in extensions that have functions
+        marked with requires_hass=True, as it assumes hass is not None.
+
+        Raises:
+            RuntimeError: If hass is not available in the environment.
+        """
+        if self.environment.hass is None:
+            raise RuntimeError(
+                "Home Assistant instance is not available. "
+                "This property should only be used in extensions with "
+                "functions marked requires_hass=True."
+            )
+        return self.environment.hass
+
     def parse(self, parser: Parser) -> Node | list[Node]:
         """Required by Jinja2 Extension base class."""
         return []

tests/helpers/template/extensions/test_base.py

@@ -0,0 +1,285 @@
+"""Test base template extension."""
+
+from __future__ import annotations
+
+import pytest
+
+from homeassistant.helpers.template import TemplateEnvironment
+from homeassistant.helpers.template.extensions.base import (
+    BaseTemplateExtension,
+    TemplateFunction,
+)
+
+
+def test_hass_property_raises_when_hass_is_none() -> None:
+    """Test that accessing hass property raises RuntimeError when hass is None."""
+    # Create an environment without hass
+    env = TemplateEnvironment(None)
+
+    # Create a simple extension
+    extension = BaseTemplateExtension(env)
+
+    # Accessing hass property should raise RuntimeError
+    with pytest.raises(
+        RuntimeError,
+        match=(
+            "Home Assistant instance is not available. "
+            "This property should only be used in extensions with "
+            "functions marked requires_hass=True."
+        ),
+    ):
+        _ = extension.hass
+
+
+def test_requires_hass_functions_not_registered_without_hass() -> None:
+    """Test that functions requiring hass are not registered when hass is None."""
+    # Create an environment without hass
+    env = TemplateEnvironment(None)
+
+    # Create a test function
+    def test_func() -> str:
+        return "test"
+
+    # Create extension with a function that requires hass
+    extension = BaseTemplateExtension(
+        env,
+        functions=[
+            TemplateFunction(
+                "test_func",
+                test_func,
+                as_global=True,
+                requires_hass=True,
+            ),
+        ],
+    )
+
+    # Function should not be registered
+    assert "test_func" not in env.globals
+    assert extension is not None  # Extension is created but function not registered
+
+
+def test_requires_hass_false_functions_registered_without_hass() -> None:
+    """Test that functions not requiring hass are registered even when hass is None."""
+    # Create an environment without hass
+    env = TemplateEnvironment(None)
+
+    # Create a test function
+    def test_func() -> str:
+        return "test"
+
+    # Create extension with a function that does not require hass
+    extension = BaseTemplateExtension(
+        env,
+        functions=[
+            TemplateFunction(
+                "test_func",
+                test_func,
+                as_global=True,
+                requires_hass=False,  # Explicitly False (default)
+            ),
+        ],
+    )
+
+    # Function should be registered
+    assert "test_func" in env.globals
+    assert extension is not None
+
+
+def test_limited_ok_functions_not_registered_in_limited_env() -> None:
+    """Test that functions with limited_ok=False are not registered in limited env."""
+    # Create a limited environment without hass
+    env = TemplateEnvironment(None, limited=True)
+
+    # Create test functions
+    def allowed_func() -> str:
+        return "allowed"
+
+    def restricted_func() -> str:
+        return "restricted"
+
+    # Create extension with both types of functions
+    extension = BaseTemplateExtension(
+        env,
+        functions=[
+            TemplateFunction(
+                "allowed_func",
+                allowed_func,
+                as_global=True,
+                limited_ok=True,  # Allowed in limited environments
+            ),
+            TemplateFunction(
+                "restricted_func",
+                restricted_func,
+                as_global=True,
+                limited_ok=False,  # Not allowed in limited environments
+            ),
+        ],
+    )
+
+    # Only the allowed function should be registered
+    assert "allowed_func" in env.globals
+    assert "restricted_func" not in env.globals
+    assert extension is not None
+
+
+def test_limited_ok_true_functions_registered_in_limited_env() -> None:
+    """Test that functions with limited_ok=True are registered in limited env."""
+    # Create a limited environment without hass
+    env = TemplateEnvironment(None, limited=True)
+
+    # Create a test function
+    def test_func() -> str:
+        return "test"
+
+    # Create extension with a function allowed in limited environments
+    extension = BaseTemplateExtension(
+        env,
+        functions=[
+            TemplateFunction(
+                "test_func",
+                test_func,
+                as_global=True,
+                limited_ok=True,  # Default is True
+            ),
+        ],
+    )
+
+    # Function should be registered
+    assert "test_func" in env.globals
+    assert extension is not None
+
+
+def test_function_registered_as_global() -> None:
+    """Test that functions can be registered as globals."""
+    env = TemplateEnvironment(None)
+
+    def test_func() -> str:
+        return "global"
+
+    extension = BaseTemplateExtension(
+        env,
+        functions=[
+            TemplateFunction(
+                "test_func",
+                test_func,
+                as_global=True,
+            ),
+        ],
+    )
+
+    # Function should be registered as a global
+    assert "test_func" in env.globals
+    assert env.globals["test_func"] is test_func
+    assert extension is not None
+
+
+def test_function_registered_as_filter() -> None:
+    """Test that functions can be registered as filters."""
+    env = TemplateEnvironment(None)
+
+    def test_filter(value: str) -> str:
+        return f"filtered_{value}"
+
+    extension = BaseTemplateExtension(
+        env,
+        functions=[
+            TemplateFunction(
+                "test_filter",
+                test_filter,
+                as_filter=True,
+            ),
+        ],
+    )
+
+    # Function should be registered as a filter
+    assert "test_filter" in env.filters
+    assert env.filters["test_filter"] is test_filter
+    # Should not be in globals since as_global=False
+    assert "test_filter" not in env.globals
+    assert extension is not None
+
+
+def test_function_registered_as_test() -> None:
+    """Test that functions can be registered as tests."""
+    env = TemplateEnvironment(None)
+
+    def test_check(value: str) -> bool:
+        return value == "test"
+
+    extension = BaseTemplateExtension(
+        env,
+        functions=[
+            TemplateFunction(
+                "test_check",
+                test_check,
+                as_test=True,
+            ),
+        ],
+    )
+
+    # Function should be registered as a test
+    assert "test_check" in env.tests
+    assert env.tests["test_check"] is test_check
+    # Should not be in globals or filters
+    assert "test_check" not in env.globals
+    assert "test_check" not in env.filters
+    assert extension is not None
+
+
+def test_function_registered_as_multiple_types() -> None:
+    """Test that functions can be registered as multiple types simultaneously."""
+    env = TemplateEnvironment(None)
+
+    def multi_func(value: str = "default") -> str:
+        return f"multi_{value}"
+
+    extension = BaseTemplateExtension(
+        env,
+        functions=[
+            TemplateFunction(
+                "multi_func",
+                multi_func,
+                as_global=True,
+                as_filter=True,
+                as_test=True,
+            ),
+        ],
+    )
+
+    # Function should be registered in all three places
+    assert "multi_func" in env.globals
+    assert env.globals["multi_func"] is multi_func
+    assert "multi_func" in env.filters
+    assert env.filters["multi_func"] is multi_func
+    assert "multi_func" in env.tests
+    assert env.tests["multi_func"] is multi_func
+    assert extension is not None
+
+
+def test_multiple_functions_registered() -> None:
+    """Test that multiple functions can be registered at once."""
+    env = TemplateEnvironment(None)
+
+    def func1() -> str:
+        return "one"
+
+    def func2() -> str:
+        return "two"
+
+    def func3() -> str:
+        return "three"
+
+    extension = BaseTemplateExtension(
+        env,
+        functions=[
+            TemplateFunction("func1", func1, as_global=True),
+            TemplateFunction("func2", func2, as_filter=True),
+            TemplateFunction("func3", func3, as_test=True),
+        ],
+    )
+
+    # All functions should be registered in their respective places
+    assert "func1" in env.globals
+    assert "func2" in env.filters
+    assert "func3" in env.tests
+    assert extension is not None