1 files changed, 76 insertions, 1 deletions

diff --git a/python/vyos/template.py b/python/vyos/template.py
index dabf53692..132f5ddde 100644
--- a/python/vyos/template.py
+++ b/python/vyos/template.py
@@ -18,7 +18,7 @@ import os
 
 from jinja2 import Environment
 from jinja2 import FileSystemLoader
-
+from jinja2 import ChainableUndefined
 from vyos.defaults import directories
 from vyos.util import chmod
 from vyos.util import chown
@@ -27,6 +27,7 @@ from vyos.util import makedir
 
 # Holds template filters registered via register_filter()
 _FILTERS = {}
+_TESTS = {}
 
 # reuse Environments with identical settings to improve performance
 @functools.lru_cache(maxsize=2)
@@ -42,8 +43,10 @@ def _get_environment(location=None):
         cache_size=100,
         loader=loc_loader,
         trim_blocks=True,
+        undefined=ChainableUndefined,
     )
     env.filters.update(_FILTERS)
+    env.tests.update(_TESTS)
     return env
 
 
@@ -67,6 +70,26 @@ def register_filter(name, func=None):
     _FILTERS[name] = func
     return func
 
+def register_test(name, func=None):
+    """Register a function to be available as test in templates under given name.
+
+    It can also be used as a decorator, see below in this module for examples.
+
+    :raise RuntimeError:
+        when trying to register a test after a template has been rendered already
+    :raise ValueError: when trying to register a name which was taken already
+    """
+    if func is None:
+        return functools.partial(register_test, name)
+    if _get_environment.cache_info().currsize:
+        raise RuntimeError(
+            "Tests can only be registered before rendering the first template"
+            )
+    if name in _TESTS:
+        raise ValueError(f"A test with name {name!r} was registered already")
+    _TESTS[name] = func
+    return func
+
 
 def render_to_string(template, content, formater=None, location=None):
     """Render a template from the template directory, raise on any errors.
@@ -566,3 +589,55 @@ def nft_intra_zone_action(zone_conf, ipv6=False):
             name = dict_search_args(intra_zone, 'firewall', fw_name)
             return f'jump {name_prefix}{name}'
     return 'return'
+
+@register_test('vyos_defined')
+def vyos_defined(value, test_value=None, var_type=None):
+    """
+    Jinja2 plugin to test if a variable is defined and not none - vyos_defined
+    will test value if defined and is not none and return true or false.
+
+    If test_value is supplied, the value must also pass == test_value to return true.
+    If var_type is supplied, the value must also be of the specified class/type
+
+    Examples:
+    1. Test if var is defined and not none:
+    {% if foo is vyos_defined %}
+    ...
+    {% endif %}
+
+    2. Test if variable is defined, not none and has value "something"
+    {% if bar is vyos_defined("something") %}
+    ...
+    {% endif %}
+
+    Parameters
+    ----------
+    value : any
+        Value to test from ansible
+    test_value : any, optional
+        Value to test in addition of defined and not none, by default None
+    var_type : ['float', 'int', 'str', 'list', 'dict', 'tuple', 'bool'], optional
+        Type or Class to test for
+
+    Returns
+    -------
+    boolean
+        True if variable matches criteria, False in other cases.
+
+    Implementation inspired and re-used from https://github.com/aristanetworks/ansible-avd/
+    """
+
+    from jinja2 import Undefined
+
+    if isinstance(value, Undefined) or value is None:
+        # Invalid value - return false
+        return False
+    elif test_value is not None and value != test_value:
+        # Valid value but not matching the optional argument
+        return False
+    elif str(var_type).lower() in ['float', 'int', 'str', 'list', 'dict', 'tuple', 'bool'] and str(var_type).lower() != type(value).__name__:
+        # Invalid class - return false
+        return False
+    else:
+        # Valid value and is matching optional argument if provided - return true
+        return True